#!/usr/bin/env python
# -*- coding: utf-8 -*-

'''
@File    : excel_handler.py
@Description : Excel处理模块，用于读取和写入Excel测试数据
'''

import os
from openpyxl import load_workbook, Workbook
from typing import Dict, List, Any, Optional, Union, Tuple

# 导入配置和日志
from config.settings import config
from utils.log_handler import get_logger


class ExcelHandler:
    """
    Excel处理类，用于读取和写入Excel测试数据
    """
    def __init__(self, file_path: str = None):
        """
        初始化Excel处理类
        
        Args:
            file_path: Excel文件路径，默认为None，使用默认测试数据文件路径
        """
        self.logger = get_logger("excel_handler")
        
        # 如果未指定文件路径，使用默认测试数据文件路径
        if file_path is None:
            self.file_path = os.path.join(config.data_dir, 'test_data.xlsx')
        else:
            self.file_path = file_path
        
        self.logger.info(f"Excel文件路径: {self.file_path}")
        
        # 检查文件是否存在
        if not os.path.exists(self.file_path):
            self.logger.warning(f"Excel文件不存在: {self.file_path}，将创建新文件")
            self._create_default_excel()
    
    def _create_default_excel(self) -> None:
        """
        创建默认的Excel测试数据文件
        """
        # 创建工作簿
        wb = Workbook()
        
        # 创建测试用例工作表
        ws = wb.active
        ws.title = "测试用例"
        
        # 添加表头
        headers = [
            "用例ID", "模块", "用例名称", "前置条件", "请求方法", "请求URL", 
            "请求头", "请求参数", "请求体", "预期状态码", "预期结果", "是否执行"
        ]
        ws.append(headers)
        
        # 添加示例数据
        example_data = [
            "TC001", "用户管理", "获取用户列表", "", "GET", "/api/users", 
            "{\"Content-Type\": \"application/json\"}", "{\"page\": 1, \"size\": 10}", "", 
            "200", "{\"code\": 0, \"message\": \"success\"}", "Y"
        ]
        ws.append(example_data)
        
        # 创建测试数据工作表
        ws2 = wb.create_sheet(title="测试数据")
        
        # 添加表头
        headers2 = ["数据ID", "用例ID", "数据名称", "数据值"]
        ws2.append(headers2)
        
        # 添加示例数据
        example_data2 = ["D001", "TC001", "用户名", "admin"]
        ws2.append(example_data2)
        
        # 保存工作簿
        # 确保目录存在
        os.makedirs(os.path.dirname(self.file_path), exist_ok=True)
        wb.save(self.file_path)
        
        self.logger.info(f"已创建默认Excel测试数据文件: {self.file_path}")
    
    def read_sheet_to_dict(self, sheet_name: str) -> List[Dict[str, Any]]:
        """
        读取指定工作表的数据，返回字典列表
        
        Args:
            sheet_name: 工作表名称
            
        Returns:
            字典列表，每个字典代表一行数据，键为表头，值为单元格值
        """
        try:
            # 使用openpyxl读取Excel
            wb = load_workbook(self.file_path, read_only=True)
            
            # 检查工作表是否存在
            if sheet_name not in wb.sheetnames:
                self.logger.error(f"工作表 {sheet_name} 不存在")
                return []
            
            ws = wb[sheet_name]
            
            # 获取表头
            headers = []
            for cell in next(ws.iter_rows(min_row=1, max_row=1, values_only=True)):
                headers.append(str(cell) if cell is not None else "")
            
            # 读取数据
            data = []
            for row in ws.iter_rows(min_row=2, values_only=True):
                row_data = {}
                for i, cell_value in enumerate(row):
                    if i < len(headers):
                        row_data[headers[i]] = cell_value
                data.append(row_data)
            
            self.logger.info(f"成功读取工作表 {sheet_name}，共 {len(data)} 行数据")
            return data
        except Exception as e:
            self.logger.error(f"读取工作表 {sheet_name} 失败: {str(e)}")
            return []
    
    # 删除依赖pandas的read_sheet_to_df方法
    
    def get_test_cases(self, filter_enabled: bool = True) -> List[Dict[str, Any]]:
        """
        获取测试用例数据
        
        Args:
            filter_enabled: 是否只返回启用的测试用例，默认为True
            
        Returns:
            测试用例数据列表
        """
        test_cases = self.read_sheet_to_dict("测试用例")
        
        # 如果需要过滤，只返回启用的测试用例
        if filter_enabled:
            test_cases = [case for case in test_cases if case.get("是否执行") == "Y"]
        
        return test_cases
    
    def get_test_data(self, case_id: str = None) -> List[Dict[str, Any]]:
        """
        获取测试数据
        
        Args:
            case_id: 测试用例ID，如果指定，则只返回该测试用例的测试数据
            
        Returns:
            测试数据列表
        """
        test_data = self.read_sheet_to_dict("测试数据")
        
        # 如果指定了测试用例ID，则只返回该测试用例的测试数据
        if case_id:
            test_data = [data for data in test_data if data.get("用例ID") == case_id]
        
        return test_data
    
    def write_result(self, case_id: str, result: str, status: str) -> None:
        """
        写入测试结果
        
        Args:
            case_id: 测试用例ID
            result: 测试结果
            status: 测试状态，如通过、失败等
        """
        try:
            # 加载工作簿
            wb = load_workbook(self.file_path)
            ws = wb["测试用例"]
            
            # 查找测试用例ID所在行
            case_id_col = 1  # 假设用例ID在第1列
            result_col = 12  # 假设测试结果在第12列
            status_col = 13  # 假设测试状态在第13列
            
            # 如果测试结果列和测试状态列不存在，则添加
            if ws.max_column < result_col:
                ws.cell(row=1, column=result_col, value="测试结果")
            if ws.max_column < status_col:
                ws.cell(row=1, column=status_col, value="测试状态")
            
            # 遍历查找测试用例ID
            for row in range(2, ws.max_row + 1):
                if ws.cell(row=row, column=case_id_col).value == case_id:
                    # 写入测试结果和状态
                    ws.cell(row=row, column=result_col, value=result)
                    ws.cell(row=row, column=status_col, value=status)
                    break
            
            # 保存工作簿
            wb.save(self.file_path)
            self.logger.info(f"成功写入测试结果，用例ID: {case_id}, 结果: {result}, 状态: {status}")
        except Exception as e:
            self.logger.error(f"写入测试结果失败: {str(e)}")
    
    def write_data_to_excel(self, data: List[Dict[str, Any]], sheet_name: str, mode: str = 'replace') -> None:
        """
        将数据写入Excel
        
        Args:
            data: 数据列表，每个元素为一个字典，代表一行数据
            sheet_name: 工作表名称
            mode: 写入模式，可选值为'replace'（替换）和'append'（追加），默认为'replace'
        """
        try:
            # 检查文件是否存在
            if not os.path.exists(self.file_path):
                # 如果文件不存在，创建新文件
                wb = Workbook()
                ws = wb.active
                ws.title = sheet_name
                
                # 如果数据不为空，写入数据
                if data:
                    # 写入表头
                    headers = list(data[0].keys())
                    ws.append(headers)
                    
                    # 写入数据
                    for row_data in data:
                        row = [row_data.get(header) for header in headers]
                        ws.append(row)
                
                # 保存工作簿
                wb.save(self.file_path)
                self.logger.info(f"成功创建Excel文件并写入工作表 {sheet_name}")
                return
            
            # 如果文件存在，根据模式写入
            wb = load_workbook(self.file_path)
            
            if mode == 'replace':
                # 替换模式：如果工作表存在，删除后重新创建；如果不存在，直接创建
                if sheet_name in wb.sheetnames:
                    # 删除工作表
                    del wb[sheet_name]
                
                # 创建新工作表
                ws = wb.create_sheet(sheet_name)
                
                # 如果数据不为空，写入数据
                if data:
                    # 写入表头
                    headers = list(data[0].keys())
                    ws.append(headers)
                    
                    # 写入数据
                    for row_data in data:
                        row = [row_data.get(header) for header in headers]
                        ws.append(row)
            elif mode == 'append':
                # 追加模式：如果工作表存在，追加数据；如果不存在，创建新工作表
                if sheet_name in wb.sheetnames:
                    # 获取工作表
                    ws = wb[sheet_name]
                    
                    # 获取表头
                    headers = []
                    for cell in next(ws.iter_rows(min_row=1, max_row=1, values_only=True)):
                        headers.append(str(cell) if cell is not None else "")
                    
                    # 追加数据
                    for row_data in data:
                        row = [row_data.get(header) for header in headers]
                        ws.append(row)
                else:
                    # 创建新工作表
                    ws = wb.create_sheet(sheet_name)
                    
                    # 如果数据不为空，写入数据
                    if data:
                        # 写入表头
                        headers = list(data[0].keys())
                        ws.append(headers)
                        
                        # 写入数据
                        for row_data in data:
                            row = [row_data.get(header) for header in headers]
                            ws.append(row)
            else:
                self.logger.error(f"不支持的写入模式: {mode}")
                return
            
            # 保存工作簿
            wb.save(self.file_path)
            self.logger.info(f"成功写入工作表 {sheet_name}，模式: {mode}")
        except Exception as e:
            self.logger.error(f"写入Excel失败: {str(e)}")


if __name__ == '__main__':
    # 测试Excel处理类
    handler = ExcelHandler()
    # 读取测试用例
    test_cases = handler.get_test_cases()
    print(f"测试用例数量: {len(test_cases)}")
    # 读取测试数据
    test_data = handler.get_test_data()
    print(f"测试数据数量: {len(test_data)}")